from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
from mobilenet import MobileNet

# 设置图片的高和宽，每一次训练所选取的样本数，迭代次数
im_height = 224
im_width = 224
batch_size = 32
# 测试集路径
validation_dir = "./data/test"
# 定义测试集图像生成器，并对图像进行预处理
# 使用图像生成器加载图片时会逐渐加载，不会一次将所有图片加载到内存中，只是将足够的图片加载到内存以便接下来的小批次使用。
# 而且可以自动缩放图像的像素值，并且可以自动生成图像的增强版本
validation_image_generator = ImageDataGenerator(rescale=1. / 255)  # 将所有数据集归一化
# 使用图像生成器从测试集路径中读取样本
val_data_gen = validation_image_generator.flow_from_directory(directory=validation_dir,  # 从测试集路径读取图片
                                                              batch_size=batch_size,  # 一次训练所选取的样本数
                                                              shuffle=False,  # 不打乱标签
                                                              target_size=(im_height, im_width),  # 图片resize到224x224大小
                                                              class_mode='categorical')  # 该参数会使得该函数返回2D的one-hot编码标签
# classes表明该分类是二分类,因为这个模型其实是支持多分类的
model = MobileNet(input_shape=(224, 224, 3), classes=2, alpha=0.25)
model.load_weights('./save/model.h5')  # 加载已训练之后的权重
# 获取数据集的类别编码，当使用图片增强器时，默认的编码即是对应的目录名称
class_indices = val_data_gen.class_indices
# 将编码和对应的类别存入字典
inverse_dict = dict((val, key) for key, val in class_indices.items())
# 加载测试图片
img = Image.open("./test/test5.jpg")
# 将图片resize到224x224大小
img = img.resize((224, 224))
# 归一化
img1 = np.array(img) / 255.
# 将图片增加一个维度，目的是匹配网络模型
img1 = (np.expand_dims(img1, 0))
# 将预测结果转化为概率值
result = np.squeeze(model.predict(img1))
predict_class = np.argmax(result)
print(inverse_dict[int(predict_class)], result[predict_class])
# 将预测的结果打印在图片上面
plt.title([inverse_dict[int(predict_class)], result[predict_class]])
# 显示图片
plt.imshow(img)
plt.show()
